#include <iostream>
#include <vector>
#include <random>

#include <opencv2/opencv.hpp>

#include "src/EKF.hpp"
#include "src/System.hpp"

using namespace std;
using namespace cv;
using std::default_random_engine;
using std::normal_distribution;

typedef Eigen::Matrix<float,4,1> StateVec;
int generatorData(vector<StateVec> &state)
{
    const float Time = 25.f,  dt = 0.1f;
    const float gravity = 9.8f;
    int length = Time/dt;
    state.reserve(length);

    default_random_engine engine;
    normal_distribution<> n(0, 1);
    for (int i = 0; i < length; ++i) {
        StateVec &lastState = state[i];
        StateVec tempState;
        tempState(0) = lastState(0) + lastState(1)*dt;
        float noiseX = n(engine);
        tempState(1) = lastState(1) + (noiseX - 0.01*lastState(1)*lastState(1))*dt;
        tempState(2) = lastState(2) + lastState(3)*dt;
        float noiseY = n(engine);
        tempState(3) = lastState(3) + (noiseY + 0.05*lastState(3)*lastState(3) - gravity)*dt;
        state.push_back(tempState);
    }

    return 0;
}

int main(int argc, char *argv[])
{
    PoseSystem poseSystem;
    StateVec initialState ;
    initialState << 0, 150, 500, 0;
    vector<StateVec> state;
    state.push_back(initialState);
    generatorData(state);

    Eigen::Matrix4f processNoise(Eigen::Vector4f(0, 0.3, 0, 0.3).asDiagonal());
    processNoise *= processNoise;
    Eigen::Matrix2f measurementNoise(Eigen::Vector2f(8,0.1).asDiagonal());
    measurementNoise *= measurementNoise;
    Eigen::Matrix4f cov = Eigen::Matrix4f::Identity();
    cov *= 10;

    Mat window(500,800,CV_8UC3,Scalar(255,255,255));
    cv::Point2f lastPoint(state[0](0),500-state[0](2));
    for (int i = 1; i < state.size(); ++i) {
        cv::Point2f currentPoint(state[i](0),500-state[i](2));
        cv::line(window,lastPoint,currentPoint,Scalar(255,0,0));
        lastPoint = currentPoint;
    }

    ExtendedKalmanFilterX<PoseSystem> ekf(poseSystem,4,0,2,processNoise,measurementNoise,initialState,cov);

    default_random_engine engine;
    normal_distribution<> n(0, 1);

     cv::Point2f lastPointEKF = cv::Point2f(state[0](0),500-state[0](2));
     cv::Point2f lastPointMeasure = cv::Point2f(state[0](0),500-state[0](2));
     for (int i = 1; i < state.size(); ++i) {
          ekf.predict();

          StateVec &currentState = state[i];
          float rho = sqrt( currentState(0)*currentState(0) + currentState(2)*currentState(2) ) + 8.*n(engine);
          float angle = atan(currentState(0)/currentState(2)) + 0.01*n(engine);

          ekf.correct(Eigen::Matrix<float,2,1>(rho,angle));

          cv::Point2f currentPointEKF(ekf.getMean()(0),500-ekf.getMean()(2));
          cv::line(window,lastPointEKF,currentPointEKF,Scalar(0,0,255));
          lastPointEKF = currentPointEKF;

          cv::Point2f currentPointMeasure(rho*sin(angle),500-rho*cos(angle));
          cv::line(window,lastPointMeasure,currentPointMeasure,Scalar(0,255,0));
          lastPointMeasure = currentPointMeasure;
     }

     cv::imshow("ResultData",window);
     cv::waitKey();

    return 0;
}
